{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7a52547",
   "metadata": {},
   "source": [
    "| [01_word_embedding/02_词向量Skip-gram模型.ipynb](https://github.com/shibing624/nlp-tutorial/blob/main/01_word_embedding/02_词向量Skip-gram模型.ipynb)  | 实现skip-gram词向量模型  |[Open In Colab](https://colab.research.google.com/github/shibing624/nlp-tutorial/blob/main/01_word_embedding/02_词向量Skip-gram模型.ipynb) |\n",
    "\n",
    "# 词向量Skip-gram模型\n",
    "\n",
    "这节训练词向量，实现Skip-gram模型。\n",
    "\n",
    "有以下几个知识点：\n",
    "- 用Skip-thought模型训练词向量\n",
    "- 学习使用PyTorch dataset和dataloader\n",
    "- 学习定义PyTorch模型\n",
    "- 学习torch.nn中常见的Module\n",
    "    - Embedding\n",
    "- 学习常见的PyTorch operations\n",
    "    - bmm\n",
    "    - logsigmoid\n",
    "- 保存和读取PyTorch模型\n",
    "\n",
    "下面复现论文[Distributed Representations of Words and Phrases and their Compositionality](http://papers.nips.cc/paper/5021-distributed-representations-of-words-and-phrases-and-their-compositionality.pdf)中训练词向量的方法，实现Skip-gram模型，使用noice contrastive sampling的目标函数，没有用论文中subsampling方法。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a9d81097",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn  # 神经网络工具箱torch.nn\n",
    "import torch.nn.functional as F  # 神经网络函数torch.nn.functional\n",
    "import torch.utils.data as tud  # Pytorch读取训练集需要用到torch.utils.data类\n",
    "\n",
    "from torch.nn.parameter import Parameter  # 参数更新和优化函数\n",
    "\n",
    "from collections import Counter  # Counter 计数器\n",
    "import numpy as np\n",
    "import random\n",
    "import math\n",
    "\n",
    "import pandas as pd\n",
    "import scipy  # SciPy是基于NumPy开发的高级模块，它提供了许多数学算法和函数的实现\n",
    "import sklearn\n",
    "from sklearn.metrics.pairwise import cosine_similarity  # 余弦相似度函数\n",
    "USE_CUDA = torch.cuda.is_available()  # 有GPU可以用,没有用CPU\n",
    "USE_CUDA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de2cda51",
   "metadata": {},
   "source": [
    "**两个模块的区别：**[torch.nn 和 torch.functional 的区别](https://blog.csdn.net/hawkcici160/article/details/80140059)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "64ee2f5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 为了保证实验结果可以复现，我们经常会把各种random seed固定在某一个值\n",
    "random.seed(42)\n",
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "if USE_CUDA:\n",
    "    torch.cuda.manual_seed(42)\n",
    "\n",
    "# 设定一些超参数\n",
    "K = 100  # number of negative samples 负样本随机采样数量\n",
    "C = 3  # nearby words threshold 指定周围三个单词进行预测\n",
    "NUM_EPOCHS = 2  # The number of epochs of training 迭代轮数，default=10\n",
    "MAX_VOCAB_SIZE = 30000  # the vocabulary size 词汇表多大\n",
    "BATCH_SIZE = 32  # the batch size 每轮迭代1个batch的数量\n",
    "LEARNING_RATE = 0.2  # the initial learning rate #学习率\n",
    "EMBEDDING_SIZE = 100  # 词向量维度\n",
    "\n",
    "\n",
    "LOG_FILE = \"word_embedding.log\"\n",
    "\n",
    "# tokenize函数，把一篇文本转化成一个个单词\n",
    "\n",
    "\n",
    "def word_tokenize(text):\n",
    "    return text.split()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90982a5a",
   "metadata": {},
   "source": [
    "- 从文本文件中读取所有的文字，通过这些文本创建一个vocabulary\n",
    "- 由于单词数量可能太大，我们只选取最常见的MAX_VOCAB_SIZE个单词\n",
    "- 我们添加一个UNK单词表示所有不常见的单词\n",
    "- 我们需要记录单词到index的mapping，以及index到单词的mapping，单词的count，单词的(normalized) frequency，以及单词总数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "59434b59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "text:  PREFACE\n",
      "\n",
      "\n",
      "SUPPOSING that Truth is a woman--what then? Is there not ground\n",
      "for suspecting that all philosophers, in so far as they have been\n",
      "dogmatists, have failed to understand women--that the terrib\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "17683"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(\"data/nietzsche.txt\", \"r\") as fin:  # 读入文件\n",
    "    text = fin.read()\n",
    "\n",
    "print('text: ', text[:200])\n",
    "\n",
    "text = [w for w in word_tokenize(text.lower())]\n",
    "# 分词，在这里类似于text.split()\n",
    "\n",
    "vocab = dict(Counter(text).most_common(MAX_VOCAB_SIZE-1))\n",
    "# 字典格式，把（MAX_VOCAB_SIZE-1）个最频繁出现的单词取出来，-1是留给不常见的单词\n",
    "\n",
    "vocab[\"<unk>\"] = len(text) - np.sum(list(vocab.values()))\n",
    "# unk表示不常见单词数=总单词数-常见单词数\n",
    "# 这里计算的到vocab[\"<unk>\"]=29999\n",
    "\n",
    "idx_to_word = [word for word in vocab.keys()]\n",
    "# 取出字典的所有单词key\n",
    "\n",
    "word_to_idx = {word: i for i, word in enumerate(idx_to_word)}\n",
    "# 取出所有单词的单词和对应的索引，索引值与单词出现次数相反，最常见单词索引为0。\n",
    "\n",
    "word_counts = np.array([count for count in vocab.values()], dtype=np.float32)\n",
    "# 所有单词的频数values\n",
    "\n",
    "word_freqs = word_counts / np.sum(word_counts)\n",
    "# 所有单词的频率\n",
    "\n",
    "word_freqs = word_freqs ** (3./4.)\n",
    "# 论文里乘以3/4次方\n",
    "\n",
    "word_freqs = word_freqs / np.sum(word_freqs)  # 用来做 negative sampling\n",
    "# 重新计算所有单词的频率\n",
    "\n",
    "VOCAB_SIZE = len(idx_to_word)  # 词汇表单词数30000=MAX_VOCAB_SIZE\n",
    "VOCAB_SIZE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e59b016",
   "metadata": {},
   "source": [
    "### 实现Dataloader\n",
    "\n",
    "一个dataloader需要以下内容：\n",
    "\n",
    "- 把所有text编码成数字，然后用subsampling预处理这些文字。\n",
    "- 保存vocabulary，单词count，normalized word frequency\n",
    "- 每个iteration sample一个中心词\n",
    "- 根据当前的中心词返回context单词\n",
    "- 根据中心词sample一些negative单词\n",
    "- 返回单词的counts\n",
    "\n",
    "这里有一个好的tutorial介绍如何使用[PyTorch dataloader](https://pytorch.org/tutorials/beginner/data_loading_tutorial.html).\n",
    "为了使用dataloader，我们需要定义以下两个function:\n",
    "\n",
    "- ```__len__``` function需要返回整个数据集中有多少个item\n",
    "- ```__getitem__``` 根据给定的index返回一个item\n",
    "\n",
    "有了dataloader之后，我们可以轻松随机打乱整个数据集，拿到一个batch的数据等等。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ab56f6ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WordEmbeddingDataset(tud.Dataset):  # tud.Dataset父类\n",
    "    def __init__(self, text, word_to_idx, idx_to_word, word_freqs, word_counts):\n",
    "        \"\"\"\n",
    "        text: a list of words, all text from the training dataset\n",
    "        word_to_idx: the dictionary from word to idx\n",
    "        idx_to_word: idx to word mapping\n",
    "        word_freq: the frequency of each word\n",
    "        word_counts: the word counts\n",
    "        \"\"\"\n",
    "        super(WordEmbeddingDataset, self).__init__()  # 初始化模型\n",
    "        self.text_encoded = [word_to_idx.get(t, VOCAB_SIZE-1) for t in text]\n",
    "        # 字典 get() 函数返回指定键的值（第一个参数），如果值不在字典中返回默认值（第二个参数）。\n",
    "        # 取出text里每个单词word_to_idx字典里对应的索引,不在字典里返回\"<unk>\"的索引\n",
    "        # \"<unk>\"的索引=29999，get括号里第二个参数应该写word_to_idx[\"<unk>\"]，不应该写VOCAB_SIZE-1，虽然数值一样。\n",
    "\n",
    "        self.text_encoded = torch.Tensor(self.text_encoded).long()\n",
    "        # 变成tensor类型，这里变成longtensor，也可以torch.LongTensor(self.text_encoded)\n",
    "\n",
    "        self.word_to_idx = word_to_idx  # 保存数据\n",
    "        self.idx_to_word = idx_to_word  # 保存数据\n",
    "        self.word_freqs = torch.Tensor(word_freqs)  # 保存数据\n",
    "        self.word_counts = torch.Tensor(word_counts)  # 保存数据\n",
    "\n",
    "    def __len__(self):  # 数据集有多少个item\n",
    "        # 魔法函数__len__\n",
    "        \"\"\"返回整个数据集（所有单词）的长度\n",
    "        \"\"\"\n",
    "        return len(self.text_encoded)  # 所有单词的总数\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # 魔法函数__getitem__，这个函数跟普通函数不一样\n",
    "        \"\"\"\n",
    "        这个function返回以下数据用于训练\n",
    "            - 中心词\n",
    "            - 这个单词附近的(positive)单词\n",
    "            - 随机采样的K个单词作为negative sample\n",
    "        \"\"\"\n",
    "        center_word = self.text_encoded[idx]\n",
    "        # print(center_word)\n",
    "        # 中心词索引\n",
    "        # 这里__getitem__函数是个迭代器，idx代表了所有的单词索引。\n",
    "\n",
    "        pos_indices = list(range(idx-C, idx)) + list(range(idx+1, idx+C+1))\n",
    "        # 周围词索引的索引，比如idx=0时。pos_indices = [-3, -2, -1, 1, 2, 3]\n",
    "\n",
    "        pos_indices = [i % len(self.text_encoded) for i in pos_indices]\n",
    "        # range(idx+1, idx+C+1)超出词汇总数时，需要特别处理，取余数\n",
    "\n",
    "        pos_words = self.text_encoded[pos_indices]\n",
    "        # 周围词索引，就是希望出现的正例单词\n",
    "        # print(pos_words)\n",
    "\n",
    "        neg_words = torch.multinomial(\n",
    "            self.word_freqs, K * pos_words.shape[0], True)\n",
    "        # 负例采样单词索引，torch.multinomial作用是对self.word_freqs做K * pos_words.shape[0]次取值，\n",
    "        # 输出的是self.word_freqs对应的下标。\n",
    "        # 取样方式采用有放回的采样，并且self.word_freqs数值越大，取样概率越大。\n",
    "        # 每个正确的单词采样K个，pos_words.shape[0]是正确单词数量\n",
    "        # print(neg_words)\n",
    "\n",
    "        return center_word, pos_words, neg_words"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e034ff69",
   "metadata": {},
   "source": [
    "创建dataset和dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d3107bb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor(5907),\n",
       " tensor([17681,     1,  5902,   705,     7,   178]),\n",
       " tensor([    5,     5,    22,     4,  1931,  1362, 15804, 14693,     9,    26,\n",
       "            40,   584, 11076, 12250,   596,  3329,    10,  5455,  3531,  1003,\n",
       "             7,  5011,   311,   300,     2,     2, 14825,    36,   181,    25,\n",
       "           153,   121,   109,   329,  3365,  3059, 11348,     2,   301,   962,\n",
       "         14926,    22,   927,   238, 16901, 11206,   635,    13,  4492,  2767,\n",
       "          5658,  1315,    12,  3218,  1602, 17616,   119,   571,  3569,   208,\n",
       "             0,  6055,   158,  3698,   475,  1037,  3336, 13463,  4919,    32,\n",
       "          2371, 16411,   301,     6,  4023,     0,  2436,  4013,    29,   272,\n",
       "          3822,  1785,  4905,    47, 13143, 16651,   121,  5867,     0,    12,\n",
       "         14703,  6160,    10,  1564, 15198,   253,  2365,  3173,     5, 16073,\n",
       "           539,   547,  9336,  2305,  5329,  2271,  1239,   548,  2706,  4298,\n",
       "           805,    57,     1,     2,    89,  7249, 14344,    12, 16038,  2603,\n",
       "          3863,    71,   195,   647,  9763,  5844,   255,  8295,   128,     0,\n",
       "         11714, 12513,  1439,   991,     8, 14310,  1427,  8572,  5760,    87,\n",
       "            11,  4180,   529,    75,   472,  2181,   655,  1465,   933,  4625,\n",
       "            24,  3519,  6972,   727,  3323,  6323,  2024,  9165,  4968,   651,\n",
       "           100, 16102,  1284,   184, 14051,  5413, 15902, 17258,  1697,    34,\n",
       "          4222, 16502,  7890, 11950,   316,    14,     2,     2,   892,  3999,\n",
       "          1653,   692,  8140,  8915,   180,  4110,  4703,     4,  6121,   315,\n",
       "          4787,  1588,     4,  2948,   385,    57,  7498,   221,  7321,    61,\n",
       "         12682,   332,     0,   130, 16285, 10784,  9868, 11457,   535,   393,\n",
       "          9578,     0,    96,   283,  2160, 12905,   859,  9896,  1094,    87,\n",
       "          8923, 13320, 11219,     9,    27,  5195,  6172,     5, 13963,  2588,\n",
       "          8183,   737,  1602,  4832, 11710,  2767,     1,  9254, 14631,  6048,\n",
       "             5,  4070,  3984,     9,  4305,    72, 14373,  5652,  1542, 16145,\n",
       "             2,  4930,    24,     4,    21,   369,   140,  1286,  1450,     1,\n",
       "          2040,  3685, 13154,  2102, 11811,    12,  1482,  5110,   128,   207,\n",
       "          2584,    18,  7787,   325,   114,    15, 13845,    23,  3632,  8596,\n",
       "          6395,    93,    70,  1344,   329,  3280,  2185,   398, 10823,     1,\n",
       "          1242,  1492,  5074,  4687,   100,  2325, 15450,     0,  4918,  8940,\n",
       "         11145,    93,   102,  1167,     2,  2190,   628,   832,  2135,  4538,\n",
       "             0,  4893, 14255,  1654,    10,   481,  2050, 13771,   628,  1583,\n",
       "          9973,  1528,     9,   286,  2472,    91,   327,  9464,   417,   471,\n",
       "             2,     6,  1372,  5844,    92,   817, 16980,    11,    16,   526,\n",
       "         10522,  4681,  5860,  1956,    29,   172,  4721,  8563,    83,  7661,\n",
       "           892,    11,     0,  5168,  1914,  4793,   258,  4636,  2944,  2577,\n",
       "         16507,  1148,  1632,   893,    21,  6700,  1832,    40, 16853,     4,\n",
       "           689,   331,   863,  9203,  9055,   174,   286,  4450,     2,   129,\n",
       "             0, 14230,     7,     5,   910,  1966,  3225,   707,   168,   716,\n",
       "          2330,     2, 13195,    12, 16241,    67,     1,  2727,  8619,   284,\n",
       "            74,   932,  7483,  2821,  1214,     2,     9,    71,   104, 15799,\n",
       "         15765,  6226, 11115,   737,     0,   843,  1563,    15,   329, 11936,\n",
       "            38,    29,  8877,  2601, 12662,   986,   158,     1,  4774,  7331,\n",
       "         11262,  2616,    17,    13,   151,  2819,     1,    48,   634,  3142,\n",
       "          7096,  3873,  1685,     7,  3349, 16217,   374, 14264,  2724,     0,\n",
       "          3042,  3319,  9573,    92,  6458, 13328,   154, 10005,   301,   410,\n",
       "          5269,  2781, 13483,  6743,  1972,  6564,    29, 13110, 14505,  4919,\n",
       "         13166,   310,   678,     5,  6434,  2950,    28,   253,    26,   286,\n",
       "             0, 13207, 13491,   130,  4819,  2584,     7,   817,  1047,   170,\n",
       "            61,     2,     7,   197,  4285,  4560,   365,    77,  2420,  7598,\n",
       "         10769, 14551, 14013,  4232,   371,     2,  2711,  2356,    66,   201,\n",
       "           202,  2437,   112,  5396,  3244, 17663, 17554,     0, 15228,     4,\n",
       "          8441,     5,   637,  4415, 13369, 12799,  2268, 15068,  1172, 15644,\n",
       "         10184,   883,    18,  1906, 12909, 12657,     8,  2544,  5784,  1835,\n",
       "         13340,  1072,   171,  4002, 14605,  1819,  2500,   621,  2918,  5626,\n",
       "           257,   187,   288,  9333,   648,   385,  1229, 15455,  1867,   693,\n",
       "         15809,    36,   240,  8767,     0,    36,   323, 15528,  5850,   111,\n",
       "          4188, 14383,    21,     2,    21,  3435,    83,  4213, 11564, 12838,\n",
       "           304,    16, 12726,   225,    81, 10057,   545, 15403,   128,    12,\n",
       "           330,  1747,  1506, 14603,  3517,  4762,    40, 13952,  9624,  5398])]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = WordEmbeddingDataset(\n",
    "    text, word_to_idx, idx_to_word, word_freqs, word_counts)\n",
    "\n",
    "# 可以把尝试打印下center_word, pos_words, neg_words看看\n",
    "list(dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b47d11c2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch.utils.data.dataloader.DataLoader at 0x7ff871d01d30>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataloader = tud.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n",
    "# notebook 下 num_workers 需要设置为0，或者在.py文件中执行。issue：https://github.com/pytorch/pytorch/issues/51344\n",
    "dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6d28bae8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 1102,   141,  1063,    35,     6,  1590,    74,    17,    11,   203,\n",
      "         2878,    11, 17405,   453,  2561,     5,   557,  1514,     5,    90,\n",
      "          217,     0,     2,   211,  1512,     0,  3294,     3,  1583,     4,\n",
      "            4, 13608]) tensor([[14499,   309,     0,     4,   203, 14500],\n",
      "        [    0,   238,     7,    24,   237,    15],\n",
      "        [ 1272,     0,  2550,   391,     3,     0],\n",
      "        [   27,   183,    18,  1469, 14183,     4],\n",
      "        [   38,  1786,    15,     5,  4944,     1],\n",
      "        [  990,   509,   203,  8385,  8386,    15],\n",
      "        [  436,    29,    82,  1503,    37,    12],\n",
      "        [    3,    67,    10,  4294,    14,    79],\n",
      "        [    2,   359,    14,    28, 12875,   501],\n",
      "        [  712,     1,    96,    29,     0, 12562],\n",
      "        [   21,    31,  7338,     3,     5,  7339],\n",
      "        [ 3301, 13091,   429,     0,  2439,   172],\n",
      "        [    3,  1496,    26,    25, 17406,   113],\n",
      "        [11140,     0,   403,     1,   269,     4],\n",
      "        [    2,     4,    50,   457,     8, 14102],\n",
      "        [    2,  1201,     8, 13234,    11,    20],\n",
      "        [ 1041,     0,   145,     1,    39,     3],\n",
      "        [   17,  1213,     0,     1,  1079,    20],\n",
      "        [    5,  6683,     3,  6684,   337,  6685],\n",
      "        [  298,     4,     0,  2309,    10,     0],\n",
      "        [  238,   652,  3473,     6,     0,   175],\n",
      "        [  765,  2979, 12349,    46,   225,    41],\n",
      "        [   40,    15,  3216,  3217,    40,    15],\n",
      "        [    5, 15614,     3,    15,    62,     8],\n",
      "        [    0,   109,  2123,     1,     0,   256],\n",
      "        [  947,     3,  5982,  5983,   134,    33],\n",
      "        [ 3514,  1009,  2519,     0,  1165,  1421],\n",
      "        [    6, 14885,   301,     0,    87,   701],\n",
      "        [    0,   574,     2,     1,     0,  1005],\n",
      "        [   36,  5071,   635,     0,  1899,    41],\n",
      "        [   15,  1623,    61,   261,     2,     4],\n",
      "        [   30,  2168,     2,    30, 13609,    14]]) tensor([[  341,  3264, 10760,  ...,  2128,     1,    16],\n",
      "        [ 1554, 14495, 10984,  ...,    16,  1069,   187],\n",
      "        [ 1059,     0,    13,  ...,  1556,  3632,   200],\n",
      "        ...,\n",
      "        [  971, 16829,   754,  ...,   115,     6,   546],\n",
      "        [13783,    21, 11532,  ..., 12575,    97,  1628],\n",
      "        [ 2433, 13519,    20,  ..., 15939,   119,  8417]])\n",
      "tensor([   28,    59, 14559,   306,   215,   176,    31,    84,   964,   626,\n",
      "         5235,   389,   703,   742, 10627,     6,  2204,     0,  4920,     0,\n",
      "        14150,   142,    11,   187,     2, 16058,   116,    39,   131,    30,\n",
      "          451,    14]) tensor([[ 1016,   106,    49,    18, 13254,    28],\n",
      "        [ 8444,     9,   226,     4,   120,    14],\n",
      "        [ 5570,    64,   861,  1647,  4472, 14560],\n",
      "        [15204,    85, 15205,   217,    34,    59],\n",
      "        [    4,   546,     3,    61,    34, 14698],\n",
      "        [   18,   766,     2,    65,   944,    19],\n",
      "        [12086,     2, 12087,   462,   402,   110],\n",
      "        [   53,    19,   199,    13,    51,  1177],\n",
      "        [   21,     4,    25,  9473,     2,  9474],\n",
      "        [ 9479,  1203,    31,     5,   203,  9480],\n",
      "        [   78,    30,   497,    24,  5076, 16726],\n",
      "        [12044,     1,     0,  5114,     2,   581],\n",
      "        [   11,    17,     0,    23,  1254,     3],\n",
      "        [    1,   640,     4,  2035, 16980, 16981],\n",
      "        [   38,    16,  1639, 10628,    30,   138],\n",
      "        [   19, 15444,   224,   768,   614,     3],\n",
      "        [ 4004,     6,    45,     8,     5,  1288],\n",
      "        [    3,  6327,  6328,   246,    11,    15],\n",
      "        [ 2813,    25,  3297,    14,   184,     3],\n",
      "        [    9,    29,   111,  8409,  8410,   186],\n",
      "        [    0,  1189,     2,  2713,     2, 14151],\n",
      "        [  121, 15152,   235,     4,     0,  3505],\n",
      "        [    1,   450,     4,   119, 15616,   722],\n",
      "        [ 1320,    33,     0,  4263,     1,   830],\n",
      "        [    2,  1264, 10616, 10617,     1,   131],\n",
      "        [ 1528,     1,  1389,   775,   299,   334],\n",
      "        [ 3863,  9391,     6,     7,     0,   277],\n",
      "        [ 2025,     1,  2587,    52,    23,   918],\n",
      "        [  299,   195,     3,  5410,   678, 13564],\n",
      "        [  195,     3,    11,    54,  5783,     5],\n",
      "        [   93, 14916,   427,    10,     0,  3513],\n",
      "        [    1,   166,   598,    16, 12641,     0]]) tensor([[  183,   461,  3086,  ...,  5245,  1430,  6511],\n",
      "        [  912,    19,   292,  ..., 16837,     2, 16398],\n",
      "        [11801,  1317,  2920,  ...,  1629,   918,   198],\n",
      "        ...,\n",
      "        [ 7612,  6927,    16,  ...,   676,     7,  7799],\n",
      "        [ 3114,  4724,  8336,  ...,   986,    32,   881],\n",
      "        [ 3541,    54,   507,  ...,    55,   131,  2220]])\n",
      "tensor([ 2222,   607,    90,     1,   695, 13435,   363,  1278,     1,    88,\n",
      "           26,   986,   158,    87,    59,     1,   452,     3, 11723,   291,\n",
      "        10960,    31,     0,     5,     5,     3,     1,     0,     3,  2482,\n",
      "         2825,  3499]) tensor([[ 4546,     1,   223,     7,     0,   740],\n",
      "        [    7,   180,   708,  6295,     2,   722],\n",
      "        [  143, 16648,     5,  2752,     7,     1],\n",
      "        [    0, 15285,   319,     5, 15286,     8],\n",
      "        [    6,     8,     5,    13,     0,  1645],\n",
      "        [   13, 13434,    68,     6,    58,    73],\n",
      "        [ 1551, 11061,   186,     3,    12,  4357],\n",
      "        [14932, 14933, 14934,     2,    91, 14935],\n",
      "        [    4,     0,  3450,     5,  2670,    21],\n",
      "        [   12,  2375,  2376,     1,   773,  1636],\n",
      "        [   65,     8,   762,     8,   120,     1],\n",
      "        [  804,     3,   113,  1438,   470,     0],\n",
      "        [   16,  2582,     2,   510,     2,  5317],\n",
      "        [   34,    59,    68,  4015,  1275,    47],\n",
      "        [    9,     6,   545,     4,  4800,  1144],\n",
      "        [ 4036,     0,   908,     0,   495,    39],\n",
      "        [   11,     6,     0,     1,  1783,    10],\n",
      "        [   10,     9,    20,    72,    14,    85],\n",
      "        [11721,    26, 11722, 11724,     2, 11725],\n",
      "        [   32,    71,   561,     1,  8128,    10],\n",
      "        [    4, 10959,     2,    22, 10961,   729],\n",
      "        [   36,   586, 16446,  2090,  3694,     0],\n",
      "        [  773,  5843,     1, 16747, 16748,   358],\n",
      "        [   39,  2834,     8,    69,  2009,     3],\n",
      "        [11300,     8,   294,   186,    83,  3015],\n",
      "        [   11,  4223,   334,   197,   406,   181],\n",
      "        [    0,   117,   212,  1574,   200,     8],\n",
      "        [   44,  1385,   121,    39,    31,     6],\n",
      "        [  696,     6,  4988,    11,    28,  2619],\n",
      "        [10392,     2,   407,     8,    19,     6],\n",
      "        [    6,   342,     8, 15426,     5,   496],\n",
      "        [   42,   102,    39,    61,   658,    27]]) tensor([[    0,  1214,    59,  ...,    48,    53,     0],\n",
      "        [ 4690,  1227,  1261,  ...,    96,  3875, 16112],\n",
      "        [10770,  1235,   244,  ...,  4169,  5666, 17432],\n",
      "        ...,\n",
      "        [  103,  6471,    55,  ...,  5828, 14311, 13359],\n",
      "        [    1,     4,     0,  ...,  3862, 12387,     5],\n",
      "        [   14,    39,     3,  ...,  7914, 17449,   250]])\n",
      "tensor([  696,  4078,    11,    20,  9289,     5,    18,    16,  5749,   654,\n",
      "        14474,    53,  1532,   168,    36,     0,     5, 12663,     8,    64,\n",
      "            2, 12677,    20,   133,     3,   113,  4541, 11996,    25,     1,\n",
      "         8675,     0]) tensor([[   87,  3549,    21,     6,  4988,     3],\n",
      "        [    0,    46,  1413,     1,  1706,     8],\n",
      "        [    0,  1434,     1,    74,   928,    54],\n",
      "        [    7,     5,   190,  5465,  4576,    27],\n",
      "        [    1,     0,  1380,    38,     0,    74],\n",
      "        [  959,   960,     8,   382,     1,  3648],\n",
      "        [ 1826,  7634,    30,    73,   300,    59],\n",
      "        [   16,  4756,   217,  3430,    27,   369],\n",
      "        [  229,  4440,     0,   363,     3,   430],\n",
      "        [  261,     2,    17,   510,    94,    67],\n",
      "        [14473,  2472,     2,    29, 14475, 14476],\n",
      "        [   22,  1787,     8,    15,    81,     0],\n",
      "        [  952,     1,     0,     1,   541,   692],\n",
      "        [    1,     0,  6432,     4,  6433,     4],\n",
      "        [14436,    47,    13,  1363,     8,     3],\n",
      "        [15106, 15107,     4,   167,     4,    11],\n",
      "        [   23,  1804,    29,  1821,    25,    71],\n",
      "        [    1,   441,     2,    30,    57,   241],\n",
      "        [14983,   148,   845,   368,     0,   341],\n",
      "        [    0,   115,  2309,    21,  4290,     8],\n",
      "        [   24,   134,  2688,  2689,     2,  6228],\n",
      "        [  359,    36, 12676,     7,    15,    52],\n",
      "        [  621,     1,    39,  6707,     3,  6708],\n",
      "        [    6,   669, 14260,   119, 14261,    35],\n",
      "        [ 1063,     6, 16524,    12,  5809,     0],\n",
      "        [ 1102,   303,   399,   540,     1,  6618],\n",
      "        [    5,   159,     1,  8945,  4541,  4542],\n",
      "        [ 2411,    18,    43,    78,    46,  5180],\n",
      "        [    2,  3358,  3432, 17554,  2598,     4],\n",
      "        [13360,     0,   450,     0,  2056,  2357],\n",
      "        [    8,     3,    16,  4475,  8676,  4476],\n",
      "        [    2,  1117,     1,  2235,     4,    32]]) tensor([[   24, 17069,  3403,  ...,     3, 12468, 15783],\n",
      "        [ 7807,  3484,  2391,  ..., 10903, 13176,     0],\n",
      "        [  131, 17600,    30,  ...,   257, 10517,  1587],\n",
      "        ...,\n",
      "        [  758,  5847,    61,  ...,  3237,  2110,  2290],\n",
      "        [    2,     0,     3,  ...,  3500,   167,  4889],\n",
      "        [12937,  7508,   154,  ...,  8079,  1058,    25]])\n"
     ]
    }
   ],
   "source": [
    "for i, (input_labels, pos_labels, neg_labels) in enumerate(dataloader):\n",
    "    print(input_labels, pos_labels, neg_labels)\n",
    "    if i > 2:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bab35705",
   "metadata": {},
   "source": [
    "torch.utils.data.DataLoader理解：https://blog.csdn.net/qq_36653505/article/details/83351808"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abc08dc0",
   "metadata": {},
   "source": [
    "### 定义PyTorch模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d00f3d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbeddingModel(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_size):\n",
    "        \"\"\" \n",
    "        初始化输出和输出embedding\n",
    "        \"\"\"\n",
    "        super(EmbeddingModel, self).__init__()\n",
    "        self.vocab_size = vocab_size  # 30000\n",
    "        self.embed_size = embed_size  # 100\n",
    "\n",
    "        initrange = 0.5 / self.embed_size\n",
    "        self.out_embed = nn.Embedding(\n",
    "            self.vocab_size, self.embed_size, sparse=False)\n",
    "        #模型输出nn.Embedding(30000, 100)\n",
    "        self.out_embed.weight.data.uniform_(-initrange, initrange)\n",
    "        # 权重初始化的一种方法\n",
    "\n",
    "        self.in_embed = nn.Embedding(\n",
    "            self.vocab_size, self.embed_size, sparse=False)\n",
    "        #模型输入nn.Embedding(30000, 100)\n",
    "        self.in_embed.weight.data.uniform_(-initrange, initrange)\n",
    "        # 权重初始化的一种方法\n",
    "\n",
    "    def forward(self, input_labels, pos_labels, neg_labels):\n",
    "        \"\"\"\n",
    "        input_labels: 中心词, [batch_size]\n",
    "        pos_labels: 中心词周围 context window 出现过的单词 [batch_size * (window_size * 2)]\n",
    "        neg_labelss: 中心词周围没有出现过的单词，从 negative sampling 得到 [batch_size, (window_size * 2 * K)]\n",
    "\n",
    "        return: loss, [batch_size]\n",
    "        \"\"\"\n",
    "\n",
    "        # input_labels是输入的标签，tud.DataLoader()返回的。相已经被分成batch了。\n",
    "        batch_size = input_labels.size(0)\n",
    "\n",
    "        input_embedding = self.in_embed(input_labels)\n",
    "        # B * embed_size\n",
    "        # 这里估计进行了运算：（128,30000）*（30000,100）= 128(B) * 100 (embed_size)\n",
    "\n",
    "        pos_embedding = self.out_embed(pos_labels)  # B * (2*C) * embed_size\n",
    "        # 同上，增加了维度(2*C)，表示一个batch有B组周围词单词，一组周围词有(2*C)个单词，每个单词有embed_size个维度。\n",
    "\n",
    "        # B * (2*C * K) * embed_size\n",
    "        neg_embedding = self.out_embed(neg_labels)\n",
    "        # 同上，增加了维度(2*C*K)\n",
    "\n",
    "        # torch.bmm()为batch间的矩阵相乘（b,n.m)*(b,m,p)=(b,n,p)\n",
    "        log_pos = torch.bmm(\n",
    "            pos_embedding, input_embedding.unsqueeze(2)).squeeze()  # B * (2*C)\n",
    "        log_neg = torch.bmm(\n",
    "            neg_embedding, -input_embedding.unsqueeze(2)).squeeze()  # B * (2*C*K)\n",
    "        # unsqueeze(2)指定位置升维，.squeeze()压缩维度。\n",
    "\n",
    "        # 下面loss计算就是论文里的公式\n",
    "        log_pos = F.logsigmoid(log_pos).sum(1)\n",
    "        log_neg = F.logsigmoid(log_neg).sum(1)  # batch_size\n",
    "        loss = log_pos + log_neg\n",
    "\n",
    "        return -loss\n",
    "\n",
    "    def input_embeddings(self):  # 取出self.in_embed数据参数\n",
    "        return self.in_embed.weight.data.cpu().numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2894059",
   "metadata": {},
   "source": [
    "定义一个模型以及把模型移动到GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "813546e0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EmbeddingModel(\n",
       "  (out_embed): Embedding(17683, 100)\n",
       "  (in_embed): Embedding(17683, 100)\n",
       ")"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = EmbeddingModel(VOCAB_SIZE, EMBEDDING_SIZE)\n",
    "# 得到model，有参数，有loss，可以优化了\n",
    "\n",
    "if USE_CUDA:\n",
    "    model = model.cuda()\n",
    "    \n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc6dfcb4",
   "metadata": {},
   "source": [
    "训练模型：\n",
    "- 模型一般需要训练若干个epoch\n",
    "- 每个epoch我们都把所有的数据分成若干个batch\n",
    "- 把每个batch的输入和输出都包装成cuda tensor\n",
    "- forward pass，通过输入的句子预测每个单词的下一个单词\n",
    "- 用模型的预测和正确的下一个单词计算cross entropy loss\n",
    "- 清空模型当前gradient\n",
    "- backward pass\n",
    "- 更新模型参数\n",
    "- 每隔一定的iteration输出模型在当前iteration的loss，以及在验证数据集上做模型的评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4bb418bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, iter: 0, loss: 100.01414489746094\n",
      "epoch: 0, iter: 100, loss: 105.59661102294922\n",
      "epoch: 0, iter: 200, loss: 91.14774322509766\n",
      "epoch: 0, iter: 300, loss: 102.48934936523438\n",
      "epoch: 0, iter: 400, loss: 88.37472534179688\n",
      "epoch: 0, iter: 500, loss: 110.75248718261719\n",
      "epoch: 0, iter: 600, loss: 125.83216857910156\n",
      "epoch: 0, iter: 700, loss: 76.50562286376953\n",
      "epoch: 0, iter: 800, loss: 158.60693359375\n",
      "epoch: 0, iter: 900, loss: 89.92433166503906\n",
      "epoch: 0, iter: 1000, loss: 75.61611938476562\n",
      "epoch: 0, iter: 1100, loss: 84.19981384277344\n",
      "epoch: 0, iter: 1200, loss: 104.96611785888672\n",
      "epoch: 0, iter: 1300, loss: 105.41595458984375\n",
      "epoch: 0, iter: 1400, loss: 102.74247741699219\n",
      "epoch: 0, iter: 1500, loss: 66.92667388916016\n",
      "epoch: 0, iter: 1600, loss: 69.28852081298828\n",
      "epoch: 0, iter: 1700, loss: 62.52679443359375\n",
      "epoch: 0, iter: 1800, loss: 88.91792297363281\n",
      "epoch: 0, iter: 1900, loss: 63.0127067565918\n",
      "epoch: 0, iter: 2000, loss: 63.874359130859375\n",
      "epoch: 0, iter: 2100, loss: 102.31259155273438\n",
      "epoch: 0, iter: 2200, loss: 59.26515579223633\n",
      "epoch: 0, iter: 2300, loss: 82.37977600097656\n",
      "epoch: 0, iter: 2400, loss: 64.3383560180664\n",
      "epoch: 0, iter: 2500, loss: 86.62437438964844\n",
      "epoch: 0, iter: 2600, loss: 71.40304565429688\n",
      "epoch: 0, iter: 2700, loss: 60.33305740356445\n",
      "epoch: 0, iter: 2800, loss: 66.86727142333984\n",
      "epoch: 0, iter: 2900, loss: 82.02959442138672\n",
      "epoch: 0, iter: 3000, loss: 57.80078125\n",
      "epoch: 1, iter: 0, loss: 35.03345489501953\n",
      "epoch: 1, iter: 100, loss: 36.674007415771484\n",
      "epoch: 1, iter: 200, loss: 35.213008880615234\n",
      "epoch: 1, iter: 300, loss: 33.516536712646484\n",
      "epoch: 1, iter: 400, loss: 34.92041778564453\n",
      "epoch: 1, iter: 500, loss: 36.32821273803711\n",
      "epoch: 1, iter: 600, loss: 33.57264709472656\n",
      "epoch: 1, iter: 700, loss: 33.5314826965332\n",
      "epoch: 1, iter: 800, loss: 35.61592102050781\n",
      "epoch: 1, iter: 900, loss: 34.53357696533203\n",
      "epoch: 1, iter: 1000, loss: 36.31746292114258\n",
      "epoch: 1, iter: 1100, loss: 33.892276763916016\n",
      "epoch: 1, iter: 1200, loss: 35.86191177368164\n",
      "epoch: 1, iter: 1300, loss: 33.94218444824219\n",
      "epoch: 1, iter: 1400, loss: 33.5980224609375\n",
      "epoch: 1, iter: 1500, loss: 35.1727180480957\n",
      "epoch: 1, iter: 1600, loss: 34.91558074951172\n",
      "epoch: 1, iter: 1700, loss: 34.41944885253906\n",
      "epoch: 1, iter: 1800, loss: 33.160667419433594\n",
      "epoch: 1, iter: 1900, loss: 35.20405578613281\n",
      "epoch: 1, iter: 2000, loss: 33.40312576293945\n",
      "epoch: 1, iter: 2100, loss: 34.19654083251953\n",
      "epoch: 1, iter: 2200, loss: 34.073768615722656\n",
      "epoch: 1, iter: 2300, loss: 33.332881927490234\n",
      "epoch: 1, iter: 2400, loss: 33.000457763671875\n",
      "epoch: 1, iter: 2500, loss: 34.13704299926758\n",
      "epoch: 1, iter: 2600, loss: 33.69338607788086\n",
      "epoch: 1, iter: 2700, loss: 33.33157730102539\n",
      "epoch: 1, iter: 2800, loss: 33.73934555053711\n",
      "epoch: 1, iter: 2900, loss: 34.01372528076172\n",
      "epoch: 1, iter: 3000, loss: 33.611106872558594\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)\n",
    "# 随机梯度下降\n",
    "\n",
    "for e in range(NUM_EPOCHS):  # 开始迭代\n",
    "    for i, (input_labels, pos_labels, neg_labels) in enumerate(dataloader):\n",
    "        # print(input_labels, pos_labels, neg_labels)\n",
    "\n",
    "        input_labels = input_labels.long()  # longtensor\n",
    "        pos_labels = pos_labels.long()\n",
    "        neg_labels = neg_labels.long()\n",
    "        if USE_CUDA:\n",
    "            input_labels = input_labels.cuda()\n",
    "            pos_labels = pos_labels.cuda()\n",
    "            neg_labels = neg_labels.cuda()\n",
    "\n",
    "        # 下面第一节课都讲过的\n",
    "        optimizer.zero_grad()  # 梯度归零\n",
    "        loss = model(input_labels, pos_labels, neg_labels).mean()\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # 打印结果。\n",
    "        if i % 100 == 0:\n",
    "            with open(LOG_FILE, \"a\") as fout:\n",
    "                fout.write(\"epoch: {}, iter: {}, loss: {}\\n\".format(\n",
    "                    e, i, loss.item()))\n",
    "                print(\"epoch: {}, iter: {}, loss: {}\".format(e, i, loss.item()))\n",
    "\n",
    "        # if i % 2000 == 0:\n",
    "        #     embedding_weights = model.input_embeddings()\n",
    "        #     sim_simlex = evaluate(\"simlex-999.txt\", embedding_weights)\n",
    "        #     sim_men = evaluate(\"men.txt\", embedding_weights)\n",
    "        #     sim_353 = evaluate(\"wordsim353.csv\", embedding_weights)\n",
    "        #     with open(LOG_FILE, \"a\") as fout:\n",
    "        #         print(\"epoch: {}, iteration: {}, simlex-999: {}, men: {}, sim353: {}, nearest to monster: {}\\n\".format(\n",
    "        #             e, i, sim_simlex, sim_men, sim_353, find_nearest(\"monster\")))\n",
    "        #         fout.write(\"epoch: {}, iteration: {}, simlex-999: {}, men: {}, sim353: {}, nearest to monster: {}\\n\".format(\n",
    "        #             e, i, sim_simlex, sim_men, sim_353, find_nearest(\"monster\")))\n",
    "\n",
    "    embedding_weights = model.input_embeddings()\n",
    "    np.save(\"embedding-{}\".format(EMBEDDING_SIZE), embedding_weights)\n",
    "    torch.save(model.state_dict(), \"embedding-{}.th\".format(EMBEDDING_SIZE))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3a3b3b22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EmbeddingModel(\n",
       "  (out_embed): Embedding(17683, 100)\n",
       "  (in_embed): Embedding(17683, 100)\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load(\"embedding-{}.th\".format(EMBEDDING_SIZE)))\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c693927b",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "下面是评估模型的代码，以及训练模型的代码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "c4dbac26",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "\n",
    "def evaluate(filename, embedding_weights):\n",
    "    if not os.path.isfile(filename):\n",
    "        return\n",
    "    if filename.endswith(\".csv\"):\n",
    "        data = pd.read_csv(filename, sep=\",\")\n",
    "    else:\n",
    "        data = pd.read_csv(filename, sep=\"\\t\")\n",
    "    print(data.head())\n",
    "    human_similarity = []\n",
    "    model_similarity = []\n",
    "    for i in data.iloc[:, 0:2].index:\n",
    "        word1, word2 = data.iloc[i, 0], data.iloc[i, 1]\n",
    "        if word1 not in word_to_idx or word2 not in word_to_idx:\n",
    "            continue\n",
    "        else:\n",
    "            word1_idx, word2_idx = word_to_idx[word1], word_to_idx[word2]\n",
    "            word1_embed, word2_embed = embedding_weights[[\n",
    "                word1_idx]], embedding_weights[[word2_idx]]\n",
    "            model_similarity.append(\n",
    "                float(sklearn.metrics.pairwise.cosine_similarity(word1_embed, word2_embed)))\n",
    "            human_similarity.append(float(data.iloc[i, 2]))\n",
    "\n",
    "    # model_similarity\n",
    "    return scipy.stats.spearmanr(human_similarity, model_similarity)\n",
    "\n",
    "\n",
    "def find_nearest(word):\n",
    "    if word not in word_to_idx:\n",
    "        return\n",
    "    index = word_to_idx.get(word, 0)\n",
    "    embedding = embedding_weights[index]\n",
    "    cos_dis = np.array([scipy.spatial.distance.cosine(e, embedding)\n",
    "                       for e in embedding_weights])\n",
    "    return [idx_to_word[i] for i in cos_dis.argsort()[:10]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "801507cd",
   "metadata": {},
   "source": [
    "## 在 Simplex-999 数据集上做评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "53784aa1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     old          new  1.58\n",
      "0  smart  intelligent  9.20\n",
      "1   hard    difficult  8.77\n",
      "2  happy     cheerful  9.55\n",
      "3   hard         easy  0.95\n",
      "4   fast        rapid  8.75\n",
      "simlex-999 SpearmanrResult(correlation=0.1552067645949596, pvalue=0.0035083019461941175)\n"
     ]
    }
   ],
   "source": [
    "embedding_weights = model.input_embeddings()\n",
    "print(\"simlex-999\", evaluate(\"data/en-simlex-999.txt\", embedding_weights))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb14bc7e",
   "metadata": {},
   "source": [
    "## 寻找nearest neighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "7d29afb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "good ['good', 'other', 'things', 'bad', 'made', 'first', 'man', 'life', 'just', 'about']\n",
      "fresh None\n",
      "monster None\n",
      "green ['green', 'judging', 'individual.=--hitherto', 'valuations\";', 'pressure', 'mixing', '\"progress', 'century,\"', 'essential.=--all', 'false,']\n",
      "like ['like', 'against', 'before', 'then', 'also', 'however,', 'ever', 'long', 'really', 'once']\n",
      "america None\n",
      "chicago None\n",
      "work ['work', 'error', 'two', 'faith', 'scientific', 'desire', 'master', 'times', 'experience', 'hand']\n",
      "computer None\n",
      "language ['language', 'mind,', 'impulse', '\"good\"', 'god,', 'difference', 'high', 'words', 'english', 'rational']\n"
     ]
    }
   ],
   "source": [
    "for word in [\"good\", \"green\", \"like\", \"america\", \"chicago\", \"work\", \"computer\", \"language\"]:\n",
    "    print(word, find_nearest(word))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1056aa58",
   "metadata": {},
   "source": [
    "## 单词之间的关系"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "c956fee0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bare?\n",
      "unmerited\n",
      "firstly\n",
      "how.\n",
      "rightly.\n",
      "natural\"\n",
      "vices--pardon\n",
      "germanise\n",
      "=reaction\n",
      "tyrannized\n",
      "atmospheric\n",
      "named\n",
      "actual\"\n",
      "uprose\n",
      "cujus\n",
      "date.\"\n",
      "grundprobleme\n",
      "suspicion.--forgive\n",
      "returning\n",
      "prevail;\n"
     ]
    }
   ],
   "source": [
    "man_idx = word_to_idx[\"man\"]\n",
    "king_idx = word_to_idx[\"king\"]\n",
    "woman_idx = word_to_idx[\"woman\"]\n",
    "embedding = embedding_weights[woman_idx] - embedding_weights[man_idx] + embedding_weights[king_idx]\n",
    "cos_dis = np.array([scipy.spatial.distance.cosine(e, embedding) for e in embedding_weights])\n",
    "for i in cos_dis.argsort()[:20]:\n",
    "    print(idx_to_word[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "3077920b",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.remove(LOG_FILE)\n",
    "os.remove('embedding-100.npy')\n",
    "os.remove('embedding-100.th')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3adeffd0",
   "metadata": {},
   "source": [
    "本节完。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81277887",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
